import rsa
import base64
import os
from rsa.key import PrivateKey,PublicKey


class MyRSA:
    def __init__(self,nbits=1024,encoding="utf-8"):
        self.nbits = nbits
        self.nbytes = (nbits+7)//8
        self.max_msg_len = self.nbytes-11
        self.pub_key, self.priv_key = rsa.newkeys(nbits)
        self.__sep = b" "
        self.__encoding = encoding

    def encrypt(self, msg: str)->bytes:
        msg = msg.encode(encoding=self.__encoding)
        ans_ls = []
        l = self.max_msg_len
        block_num = (len(msg)+l)//l
        for i in range(block_num):
            encrypted_block = rsa.encrypt(msg[i*l:(i+1)*l],self.pub_key)
            b64_block = base64.b64encode(encrypted_block)
            ans_ls.append(b64_block)
        ans = self.__sep.join(ans_ls)
        return ans
    

    def decrypt(self, cipertext: bytes)->str:
        ans_ls = []
        b64_blocks = cipertext.split(self.__sep)
        for b64_block in b64_blocks:
            encrypted_block = base64.b64decode(b64_block)
            block = rsa.decrypt(encrypted_block,self.priv_key)
            ans_ls.append(block)
        ans = b"".join(ans_ls)
        ans =ans.decode(encoding=self.__encoding)
        return ans

